import torch
from .base_model import BaseModel
from . import networks


class Pix2PixModel(BaseModel):
    @staticmethod
    def modify_commandline_options(parser, is_train=True):
        parser.set_defaults(norm='batch', netG='unet_256', dataset_mode='aligned')
        if is_train:
            # 训练时候使用交叉熵作为生成数据的判断，同时增加了一个L1正则化选项
            parser.set_defaults(pool_size=0, gan_mode='vanilla')
            parser.add_argument('--lambda_L1', type=float, default=100.0, help='weight for L1 loss')

        return parser

    def __init__(self, opt, isTrain):
        BaseModel.__init__(self, opt)

        self.isTrain = isTrain
        # 定义生成器网络
        self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm,
                                      not opt.no_dropout, opt.init_type, opt.init_gain)
        # 训练时候定义一个判别器。
        if self.isTrain:  # define a discriminator; conditional GANs need to take both input and output images; Therefore, #channels for D is input_nc + output_nc
            self.netD = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
                                          opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain)
        # 训练时候定义了loss 函数，分为生成LOSS和L1正则化LOSS两种。
        if self.isTrain:
            # define loss functions
            # vanilla| lsgan | wgangp 默认lsgan
            # GANLoss 也被定义成为了一个model
            # print("opt.gan_mode", opt.gan_mode)
            self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device)
            self.criterionL1 = torch.nn.L1Loss()
            # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
            # 定义两个优化器，一个负责优化生成器，一个负责优化判别器。
            self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
            self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
            self.optimizers.append(self.optimizer_G)
            self.optimizers.append(self.optimizer_D)

    def set_input(self, input):
        AtoB = self.opt.direction == 'AtoB'
        self.real_A = input['A' if AtoB else 'B'].to(self.device)
        self.real_B = input['B' if AtoB else 'A'].to(self.device)
        self.real_A = torch.reshape(self.real_A, (1, 3, 256, 256))
        self.real_B = torch.reshape(self.real_B, (1, 3, 256, 256))
        self.image_paths = input['A_paths' if AtoB else 'B_paths']

    def forward(self):
        self.fake_B = self.netG(self.real_A)  # G(A)

    def backward_D(self):
        """Calculate GAN loss for the discriminator"""
        # Fake; stop backprop to the generator by detaching fake_B
        # 将生成的B图和A图合并作为输入
        fake_AB = torch.cat((self.real_A, self.fake_B),
                            1)  # we use conditional GANs; we need to feed both input and output to the discriminator
        # 将fack_AB的梯度锁住不更新，并送入判别器，之所以锁住不更新该参数，是因为这里把B生成的图像当作静态图像数据来处理。
        pred_fake = self.netD(fake_AB.detach())
        # 计算从A和生成的B合并之后的图片经过判别器后被认为假的交叉熵损失函数
        self.loss_D_fake = self.criterionGAN(pred_fake, False)
        # Real
        real_AB = torch.cat((self.real_A, self.real_B), 1)
        pred_real = self.netD(real_AB)
        # 计算将真实A和真是B合并之后的图片经过判别器后被认为真的交叉熵损失函数
        self.loss_D_real = self.criterionGAN(pred_real, True)
        # combine loss and calculate gradients
        # 将两个交叉熵损失函数加和并乘以0.5，用于做判别器的参数梯度下降。关键一步，锁住不更新fakeB数据。
        self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5
        self.loss_D.backward()

    def backward_G(self):
        """Calculate GAN and L1 loss for the generator"""
        # First, G(A) should fake the discriminator
        # 将真实的A和生成的B通过判别器计算和真图像的交叉熵损失
        fake_AB = torch.cat((self.real_A, self.fake_B), 1)
        pred_fake = self.netD(fake_AB)
        self.loss_G_GAN = self.criterionGAN(pred_fake, True)
        # Second, G(A) = B
        # 将生成的B和真实的B进行每个元素之间的平均绝对误差（MAE）计算。
        self.loss_G_L1 = self.criterionL1(self.fake_B, self.real_B) * 100.0
        # combine loss and calculate gradients
        # 将两者同时相加并更新参数
        self.loss_G = self.loss_G_GAN + self.loss_G_L1
        self.loss_G.backward()

    def optimize_parameters(self):
        self.forward()  # compute fake images: G(A)
        # update D
        self.set_requires_grad(self.netD, True)  # enable backprop for D
        self.optimizer_D.zero_grad()  # set D's gradients to zero
        self.backward_D()  # calculate gradients for D
        self.optimizer_D.step()  # update D's weights
        # update G
        self.set_requires_grad(self.netD, False)  # D requires no gradients when optimizing G
        self.optimizer_G.zero_grad()  # set G's gradients to zero
        self.backward_G()  # calculate graidents for G
        self.optimizer_G.step()  # udpate G's weights
